package com.chinanet.config;

import org.apache.catalina.connector.Connector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.web.embedded.tomcat.TomcatConnectorCustomizer;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


public class GracefulShutdown implements TomcatConnectorCustomizer,ApplicationListener<ContextClosedEvent> {
    private static final Logger log =LoggerFactory.getLogger(GracefulShutdown.class);

    private static final int TIMEOUT=60;

    private volatile Connector connector;

    @Override
    public void customize(Connector connector) {
        this.connector=connector;
    }

    @Override
    public void onApplicationEvent(ContextClosedEvent contextClosedEvent) {
        this.connector.pause();
        Executor executor =this.connector.getProtocolHandler().getExecutor();
        if(executor instanceof ThreadPoolExecutor){
            ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) executor;
            threadPoolExecutor.shutdown();
            try {
                if(!threadPoolExecutor.awaitTermination(TIMEOUT,TimeUnit.SECONDS)){
                    log.warn("Tomcat thread pool did not shut down gracefully within "+TIMEOUT+" seconds. Proceeding with forceful shutdown");
                    threadPoolExecutor.shutdownNow();
                    if(!threadPoolExecutor.awaitTermination(TIMEOUT,TimeUnit.SECONDS)){
                        log.error("Tomcat thread pool did not terminate");
                    }
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }
}
